/*
 * @Author: LibraStalker 1299640057@qq.com
 * @Date: 2023-01-27 23:44:51
 * @FilePath: landmark68.ts
 * @Description: 人脸关键点检测
 */
import {Runner, interfaces} from '@paddlejs/paddlejs-core';
import {ModelSDK} from "./base";
import {FeedShape, Location, ModelConfig, NaturalSize, Coordinate} from "../commons/interface";

/**
 * 人脸关键点检测
 */
export class Landmark68 extends ModelSDK {
    readonly defaultFeedShape: FeedShape = {
        fc: 1,
        fw: 60,
        fh: 60
    };

    // 这里的平均值和标准差需要根据图片自动计算，所以不需要配置
    readonly defaultModelConfig: ModelConfig = {
        modelPath: '../../model/face_landmark_localization/model.json',
        feedShape: this.defaultFeedShape
    };

    inputSize = {} as NaturalSize // 原始图片的尺寸
    scale: number; // feedShape / inputSize （指长度或宽度的放缩比例）
    dx: number; // 左侧黑边的宽度
    dy: number; // 上侧黑边的高度

    constructor(modelConfig?: ModelConfig) {
        super(modelConfig);
        super.__init__();
    }

    async init(): Promise<any> {
        this.runner = new Runner(this.modelConfig);
        // Runner.init的过程中会调用MediaProcessor.allReshapeToRGB，但这个函数是根据输入的
        // 平均值和标准差来进行数据的标准化的，无法自动计算平均值和标准差，而且无法转灰度，所以这里
        // 将其更改为将输入数据转为灰度后自动计算平均值和标准差，再进行标准化
        this.runner.mediaProcessor.allReshapeToRGB = (imageData, opt): Float32Array => {
            const {
                targetShape,
                bgr,
                normalizeType = 0
            } = opt;
            const [, c, h, w] = targetShape;
            const data = imageData.data || imageData; //取数据
            const gray_array = this.RGBAtoGRAY(data); // 转灰度
            const mean = this.calculate_mean(gray_array); // 平均值
            const std = this.calculate_std(gray_array); // 标准差
            const result = new Float32Array(gray_array.length);
            for (let i = 0; i < result.length; i++) {
                result[i] = gray_array[i] - mean; // 平均值置为0
                result[i] /= std; // 标准差置为1
            }
            return result;
        };
        return this.runner.init();
    }

    /**
     * 从仅包含人脸的图片中检测一张人脸的关键点
     * @param input 原始图片中人脸检测框框出的部分，即仅包含人脸的图片，可以通过createImageBitmap()获取
     * @param base_location 人脸图片在原始图片中的位置
     * @param base_location.left 人脸图片的左边界与原始图片的左边界的像素距离，一般为非负数
     * @param base_location.top 人脸图片的上边界与原始图片的上边界的像素距离，一般为非负数
     * @param base_location.width 人脸图片的像素宽度，正数
     * @param base_location.height 人脸图片的像素高度，正数
     * @return 坐标数组，Coordinate中的x为关键点距离原始图片左边界的像素距离，y为关键点距离原始图片上边界的像素距离
     */
    async detect(
        input: ImageBitmap,
        base_location: Location = {
            left: 0, top: 0, width: 0, height: 0
        }
        // opts?
    ): Promise<Coordinate[]> {
        this.inputSize.naturalWidth = input.width;
        this.inputSize.naturalHeight = input.height;

        const inputFeed = this.preprocess(input); // 预处理
        console.debug(inputFeed);
        const dataOut: number[] = await this.runner.predictWithFeed(inputFeed); // 预测
        console.debug(dataOut);
        return this.postprocess(dataOut, base_location); // 后处理
    }

    /**
     * 预处理，将输入的图像数据按比例缩放到feedShape，若长宽比与feedShape不同则多余的部分留空，
     * 再转为灰度，并转为CHW格式，即1 * 60 * 60
     * @param input 原始图片中人脸检测框框出的部分，即仅包含人脸的图片，可以通过createImageBitmap()获取
     * @return InputFeed[] 可以直接传给Runner.predictWithFeed
     */
    preprocess(input: ImageBitmap): interfaces.InputFeed[] {
        return this.runner.mediaProcessor.process(
            input,
            Object.assign({}, this.modelConfig),
            this.feedShape
        );
    }

    /**
     * 后处理，将模型输出的长度为(2 * 关键点个数)的比例坐标数组转换为关键点相对于原始图片的像素坐标数组。
     *
     * 由于用到了实例的属性，所以不是异步安全的。
     * @param dataOut 模型的输出，长度为(2 * 关键点个数)的一维数组
     * @param base_location 人脸图片在原始图片中的位置
     * @param base_location.left 人脸图片的左边界与原始图片的左边界的像素距离，一般为非负数
     * @param base_location.top 人脸图片的上边界与原始图片的上边界的距离，一般为非负数
     * @param base_location.width 人脸图片的像素宽度，正数
     * @param base_location.height 人脸图片的像素高度，正数
     * @return 坐标数组，Coordinate中的x为关键点距离原始图片左边界的像素距离，y为关键点距离原始图片上边界的像素距离
     */
    postprocess(dataOut: number[], base_location: Location): Coordinate[] {
        const result = [];
        const {naturalWidth, naturalHeight} = this.inputSize; // 原始图片的宽高
        const {fw, fh} = this.feedShape; // 给模型输入的宽高

        if (naturalWidth > naturalHeight) {
            this.scale = fw / naturalWidth;
        } else {
            this.scale = fh / naturalHeight;
        }
        this.dx = (fw - naturalWidth * this.scale) / 2; // 左侧白边的宽度
        this.dy = (fh - naturalHeight * this.scale) / 2; // 上侧白边的高度
        for (let i = 0; i < dataOut.length; i += 2) {
            result.push(this.transformData(dataOut[i], dataOut[i + 1], base_location));
        }
        return result;
    }

    /**
     * 将模型输出的比例坐标转换为关键点相对于原始图片的像素坐标。
     *
     * 由于用到了实例的属性，所以不是异步安全的。
     * @param x 关键点距离inputFeed的图像的左边界的距离 / inputFeed的图像的宽度
     * @param y 关键点距离inputFeed的图像的上边界的距离 / inputFeed的图像的高度
     * @param base_location 人脸图片在原始图片中的位置
     * @param base_location.left 人脸图片的左边界与原始图片的左边界的像素距离，一般为非负数
     * @param base_location.top 人脸图片的上边界与原始图片的上边界的距离，一般为非负数
     * @param base_location.width 人脸图片的像素宽度，正数
     * @param base_location.height 人脸图片的像素高度，正数
     * @return 关键点的绝对坐标，Coordinate中的x为关键点距离原始图片左边界的像素距离，y为关键点距离原始图片上边界的像素距离
     */
    transformData(x: number, y: number, base_location: Location): Coordinate {
        const {fw, fh} = this.feedShape;
        return {
            x: (x * fw - this.dx) / this.scale + base_location.left,
            y: (y * fh - this.dy) / this.scale + base_location.top
        };
    }
}